# 示例：如何使用内存组件

import numpy as np
import torch
from torch.utils.data import DataLoader

from learners.components.__further__.episodic_memory import EpisodicMemory
from learners.components.__further__.memory import (
    HerdingExemplarSelector,
    FixedMemoryStrategy,
)
from learners.components.__further__.memory_integration import refactor_learner_memory


def example_standalone_memory_usage():
    """使用独立的内存组件示例"""

    # 1. 创建一个特征提取器（实际中可能是模型的某个部分）
    def feature_extractor(inputs):
        # 这里简单模拟提取特征，实际中应该是调用模型
        return torch.randn(inputs.size(0), 512)  # 假设特征维度为512

    # 2. 创建内存组件
    memory = EpisodicMemory(
        feature_extractor=feature_extractor,
        feature_dim=512,
        memory_size=2000,  # 总的内存大小限制
        memory_per_class=20,  # 每类保留的样本数
        fixed_memory=True,  # 使用固定内存策略
        total_classes=100,  # 总类别数
        device="cuda",
        batch_size=64,
    )

    # 3. 使用内存组件
    # ... 在训练过程中使用 memory.build_rehearsal_memory() 来构建记忆
    # ... 使用 memory.get_memory() 来获取记忆中的数据

    print("创建了具有以下配置的内存组件:")
    print(f"  - 特征维度: {memory.feature_dim}")
    print(f"  - 每类样本数: {memory.memory_per_class}")
    print(f"  - 固定内存策略: {memory.fixed_memory}")


def example_integrate_with_learner(learner):
    """将内存组件集成到现有的学习器中"""

    # 重构学习器以使用组件化的内存系统
    refactor_learner_memory(learner)

    print("已将学习器重构为使用组件化内存系统")
    print("现在可以正常使用学习器的内存相关方法，但它们会使用新的内存组件")


def custom_memory_example():
    """使用自定义的内存选择策略示例"""

    # 创建一个特征提取器
    def feature_extractor(inputs):
        return torch.randn(inputs.size(0), 512)

    # 创建自定义的内存策略
    custom_strategy = FixedMemoryStrategy(memory_per_class=30)

    # 创建内存组件
    memory = EpisodicMemory(
        feature_extractor=feature_extractor,
        feature_dim=512,
        memory_size=2000,
        memory_per_class=30,
        fixed_memory=True,
        total_classes=100,
        device="cuda",
    )

    # 手动替换内存策略
    memory.memory_strategy = custom_strategy

    print("创建了带有自定义内存策略的内存组件")


if __name__ == "__main__":
    example_standalone_memory_usage()
    # 若要运行其他示例，请取消注释以下行
    # example_integrate_with_learner(your_learner)
    # custom_memory_example()
